package datastructure.tree;

import java.util.Comparator;

import static datastructure.tree.RedBlackNode.BLACK;
import static datastructure.tree.RedBlackNode.RED;

/**
 * @Description: 红黑树
 * @Author Ammar
 * @Create 2023/4/17 07:12
 */
public class RedBlackTree<E> extends BalaceBinarySearchTree<E> {

    public RedBlackTree() {
        this(null);
    }

    public RedBlackTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 添加节点一共有12种情况
     * 首先默认添加的都是红色节点，这样可以让红黑树的5个性质满足4个，不会满足的性质为：Red节点的子节点必须是Black节点
     * 1. 当 parent = Black 时，不会违反性质4，所以不用处理
     * 2. 当 parent = Red 时，有 8 种情况不满足性质4
     * 2.1 其中 uncle = Red 时，属于B树种的上溢，需要将 parent 和 uncle 都染成 Black，grand 向上合并并且染成 Red，当做是新添加的节点进行处理(递归)
     * 2.2 其中 uncle = Black 时，uncle 不是 Red，有可能是 LL\RR\LR\RL 四种情况
     *
     * @param node
     */
    @Override
    protected void afterAdd(TreeNode<E> node) {
        TreeNode<E> parent = node.parent;

        // 添加的是根节点
        if (parent == null) {
            black(node);
            return;
        }
        // 如果父节点是黑色，直接返回
        if (isBlack(parent)) return;

        // uncle节点
        TreeNode<E> uncle = parent.sibling();
        TreeNode<E> grand = red(parent.parent);

        // 如果uncle节点是红色，上溢
        // 不管是哪种情况，都要将 grand 节点染成红色，parent 和 uncle 节点染成黑色
        if (isRed(uncle)) {
            black(parent);
            black(uncle);
            // 把祖父节点当做是新添加的节点
//            TreeNode<E> red = red(parent.parent);
            afterAdd(grand);
            return;
        }
        if (parent.isLeftChild()) { // L
            // red(grand);
            if (node.isLeftChild()) { // LL
                black(parent);
            } else { // LR
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        } else { // R
            // red(grand);
            if (node.isLeftChild()) { // RL
                black(node);
                rotateRight(parent);
            } else { // RR
                black(parent);
            }
            rotateLeft(grand);
        }
    }


    @Override
    protected void afterRemove(TreeNode<E> node, TreeNode<E> replacement) {
        // 如果要删除的节点是红色，什么也不做
        if (isRed(node)) return;

        // 如果删除的是黑色且度为1的节点，用子节点替代，子节点必然是红色，将其染色黑色即可
        if (isRed(replacement)) {
            black(replacement);
            return;
        }

        TreeNode<E> parent = node.parent;
        // 此时，红黑树只有一个根节点，
        // 如果还有其他节点，删除根节点时，实际上删除的是前驱或者后继节点
        if (parent == null) return;

        // 1.程序来到这里，说明删除的是度为0的节点
        // 2.按照4阶B树的性质，会产生下溢
        // 3.如果兄弟节点有 RED 子节点，则借过来，所以要获取兄弟节点，
        // 4.二叉树删除叶子节点后，parent.left/right 已经指向null了，但是node.parent还未清除
        //   根据红黑树的性质，黑色叶子节点必然有BLACK兄弟节点，否则不满足性质5
        //   因此，在节点删除后，判断 parent.left/right 哪边为空，就可以知道删除的节点在 parent 的哪边了
        //   有一种特殊情况，就是在度为0的BLACK节点被删除后，parent下溢，
        //   然而，parent没有RED的子节点，parent下溢以后导致grand下溢，又会调用 afterRemove，
        //   此时parent.left/right都不为空，因为grand只是下溢没有被删除，
        //   这时需要判断grand是grand.parent的left or right，此时才能找到grand的sibling
        boolean left = parent.left == null || node.isLeftChild();
        TreeNode<E> sibling = left ? parent.right : parent.left;

        // node在parent左边
        if (left) {
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateLeft(parent);
                sibling = parent.right;
            }
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent, null);
                }
            }
            else {
                if (isBlack(sibling.right)) {
                    rotateRight(sibling);
                    sibling = parent.right;
                }
                color(sibling, colorOf(parent));
                black(sibling.right);
                black(parent);
                rotateLeft(parent);
            }
        }
        // node在parent右边
        else {
            if (isRed(sibling)) {
                black(sibling);
                red(parent);
                rotateRight(parent);
                // 更换 sibling
                sibling = parent.left;
            }
            // sibling 左右子节点全是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // parent 是否有 RED true 表示没有 RED
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if (parentBlack) {
                    afterRemove(parent, null);
                }
            }
            // sibling 至少有一个红色子节点，将子节点借过来，需要旋转
            // 一共三种情况，一 sibling.right = RED LR
            // 二 sibling.left = RED LL
            // 三 sibling.left and sibling.right = RED LL
            else {
                // LR情况
                if (isBlack(sibling.left)) {
                    rotateLeft(sibling);
                    sibling = parent.left;
                }

                // 由于旋转之后再获取对应的节点会有问题，因此先染色，再旋转
                color(sibling, colorOf(parent));
                black(sibling.left);
                black(parent);
                rotateRight(parent);
            }
        }
    }

    // 给节点染色
    private TreeNode<E> color(TreeNode<E> node, boolean color) {
        if (node == null) return node;
        ((RedBlackNode<E>) node).color = color;
        return node;
    }

    // 将节点染成红色
    private TreeNode<E> red(TreeNode<E> node) {
        return color(node, RED);
    }

    // 将节点染成黑色
    private TreeNode<E> black(TreeNode<E> node) {
        return color(node, BLACK);
    }

    // 返回节点的颜色
    private boolean colorOf(TreeNode<E> node) {
        return node == null ? BLACK : ((RedBlackNode<E>) node).color;
    }

    // 判断节点是否为黑色
    private boolean isBlack(TreeNode<E> node) {
        return colorOf(node) == BLACK;
    }

    // 判断节点是否为红色
    private boolean isRed(TreeNode<E> node) {
        return colorOf(node) == RED;
    }

    @Override
    protected TreeNode<E> createNode(E val, TreeNode<E> parent) {
        return new RedBlackNode<>(val, parent);
    }
}
